import os
import sys
import cv2
import glob
import time
import math
import ntpath
import shutil
import corecode.api as corecode
from corecode.textbox import GiP
import numpy as np
import matplotlib.pyplot as plt
np.set_printoptions(threshold=32, edgeitems=10)
try:
   import cPickle as pickle
except:
   import pickle

def path_leaf(path):
    head, tail = ntpath.split(path)
    return tail or ntpath.basename(head)

def storeEm(imagefile, naS=None):
    cv_i = cv2.imread(imagefile)
    
    if imagefile is None:
            print ('Not an image')
            sys.exit(1)
    if str(type(cv_i)) == "<class 'NoneType'>":
            print ('Not an image')
            sys.exit(1)
    if cv_i.shape[1] < 256:
            print ('This script is for getting storing high-resolution face encodings')
            print ('Please select a higher-resolution face.')
            sys.exit(1)
    img1 = corecode.load_image_file(imagefile)
    fen0 = corecode.face_encodings( img1)
    if len(fen0) > 1:
            print ('This script is for storing the embeddings from a single face only')
            print ('Please select a single face.')
            sys.exit(1)
    fen1 = corecode.face_encodings( img1)[0]
    print("naS: " + naS)
    if (' ') in naS:

       p_fn, p_ln = naS.split(' ')
       f_name = str(p_fn)
       l_name = str(p_ln)
       name = ("{}.{}".format(f_name.lower(),l_name.lower()))
    elif naS == None:
       name = path_leaf(imagefile)
    else:
       name = naS


    def makeFolder(n):
        folder = ((os.getcwd()) + "/embeddings/" + n + "/")
        if not os.path.exists(folder): os.makedirs(folder)
        return folder
    
    e_f = makeFolder(name)
        
    def  filecount(DIR): 
         return len([name for name in os.listdir(DIR) if os.path.isfile(os.path.join(DIR, name))])
    
    cInt = filecount(e_f)
    
    if cInt < 2:
       fID = str(cInt + 1).zfill(4)
       with open(e_f + "/{}_mainEm.{}_".format(str(fID),name) + ".pickle",'wb') as pf: 
            pickle.dump(fen1, pf)
       fex = str(imagefile).split('.')[-1]
       shutil.copy(imagefile, e_f + "/{}_mainPk.{}_.".format(str(fID),name) + fex)

    else:
    
       tInt = math.ceil(cInt/2)
    
       print ("AdId #{}".format(tInt))
    
       fID = str(tInt + 1).zfill(4)
       with open(e_f + "/{}_ad_mainEm.{}_".format(str(fID),name) + ".pickle",'wb') as pf: 
            pickle.dump(fen1, pf)
       fex = str(imagefile).split('.')[-1]
       shutil.copy(imagefile, e_f + "/{}_ad_mainPk.{}_.".format(str(fID),name) + fex)
    


def main():
    path = sys.argv[1]
    types = ('*.jpg','*.jpeg','*.png','*.gif',)
    f_g = []
    for files in types:
        f_g.append( glob.glob(os.path.join(path, files)) )
    for i in f_g:
        for r in i:
            print(r)
            plt.figure()
            i_f = corecode.load_image_file(r)

            plt.imshow(i_f)
            plt.show(block=False)
            personname = GiP()
            storeEm(r, personname)

            plt.close()





main()








